/**
 * 
 */
package edu.umd.clip.lm.programs;

import java.io.File;
import edu.berkeley.nlp.util.*;

import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.training.*;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class Interpolate {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-jobs", required = false, usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-lm", required = false, usage = "LM ID to train (default: " + LanguageModel.PRIMARY_LM_ID + ")")
		public String lm = LanguageModel.PRIMARY_LM_ID;        
        @Option(name = "-backoff", required = false, 
        		usage = "method to compute backoff weight ("+BACKOFF_ONE+"|"+BACKOFF_ZERO+"|"+BACKOFF_HEURISTIC+") (default: "+BACKOFF_HEURISTIC+")")
		public String backoff = BACKOFF_HEURISTIC;        
        @Option(name = "-datadir", required = false, usage = "the directory for the temporary files")
		public String datadir = null;        
        @Option(name = "-jerboa", required = false, usage = "use Jerboa storage (default: false)")
		public boolean useJerboa = false;        
        @Option(name = "-compact", required = false, usage = "use Compact storage (default: false)")
		public boolean useCompact = false;        
        @Option(name = "-bdb", required = false, usage = "use Berkeley DB storage (default: false)")
		public boolean useBDB = false;        
	}
	
	private static final String BACKOFF_ONE = "1";
	private static final String BACKOFF_ZERO = "0";
	private static final String BACKOFF_HEURISTIC = "heuristic";
	
	public static void main(String[] args) throws Exception {
        OptionParser optParser = new OptionParser(Options.class);
        final Options opts = (Options) optParser.parse(args, true);

        LMDecodingOptions lmOpts = new LMDecodingOptions();
        lmOpts.config = opts.config;
        lmOpts.jobs = opts.jobs;
        if (opts.lm != null) lmOpts.lm = opts.lm;
        
        if (opts.useJerboa) lmOpts.storage = LMDecodingOptions.Storage.JERBOA;
        if (opts.useBDB) lmOpts.storage = LMDecodingOptions.Storage.BDB;
        if (opts.useCompact) lmOpts.storage = LMDecodingOptions.Storage.COMPACT;

        lmOpts.dontOpenCurrentLMStorage = true;
        
        LanguageModel.initDecoding(lmOpts);

        Experiment experiment = Experiment.getInstance();
		
		LanguageModel lm  = experiment.getLM(opts.lm);

		File tmpDir = opts.datadir == null ? new File("populate-" + lm.getId()) : new File(opts.datadir);
		if (!tmpDir.isDirectory()) {
			tmpDir.mkdirs();
		}
		
		BackoffType backoff;
		
        if (opts.backoff.equals(BACKOFF_ONE)) {
        	backoff = BackoffType.ONE;
        } else if (opts.backoff.equals(BACKOFF_ZERO)) {
        	backoff = BackoffType.ZERO;
        } else {
        	backoff = BackoffType.HEURISTICS;
        }

        BackoffInterpolation interpolator = new BackoffInterpolation(lm);
        interpolator.setBackoff(backoff);

        interpolator.initialize(tmpDir);
        interpolator.interpolate();
        
		lm.saveHistoryTree();

		LanguageModel boLM = experiment.getLM(lm.getBackoffLM());
		if (boLM != null) {
			boLM.getDecoder().getStorage().closeAll();
		}
	}
}
